package middleware

import (
	"fmt"
	"github.com/gin-gonic/gin"
	"net/http"
	"strings"
	"sync"
	"testGin/tool"
	"time"
)

var limiter = NewRateLimiter()

type RateLimiter struct {
	bucket map[string]*TokenBucket
	mutex  sync.Mutex
}

type TokenBucket struct {
	rate       float64   // 速率，单位：令牌/秒
	capacity   float64   // 令牌桶容量
	tokens     float64   // 当前令牌数量
	lastUpdate time.Time // 上次更新时间
}

func NewRateLimiter() *RateLimiter {
	return &RateLimiter{
		bucket: make(map[string]*TokenBucket),
	}
}

func (rl *RateLimiter) AllowIP(ip string) bool {
	rl.mutex.Lock()
	defer rl.mutex.Unlock()

	bucket, exists := rl.bucket[ip]
	if !exists {
		// 初始化令牌桶
		bucket = &TokenBucket{
			rate:       3, // 每秒生成5个令牌
			capacity:   3, // 令牌桶容量为5个
			tokens:     3, // 初始时令牌桶为满的状态
			lastUpdate: time.Now(),
		}
		rl.bucket[ip] = bucket
	}
	// 计算时间间隔，并根据速率生成令牌
	now := time.Now()
	elapsed := now.Sub(bucket.lastUpdate).Seconds()
	tokensToAdd := elapsed * bucket.rate

	// 更新令牌桶状态
	if tokensToAdd > 0 {
		bucket.tokens = bucket.tokens + tokensToAdd
		if bucket.tokens > bucket.capacity {
			bucket.tokens = bucket.capacity
		}
		bucket.lastUpdate = now
	}
	// 检查令牌数量是否足够
	if bucket.tokens >= 1 {
		bucket.tokens--
		return true
	}
	return false
}

func QPSLimit() gin.HandlerFunc {
	return func(c *gin.Context) {
		ip := c.Request.Header.Get("X-Real-IP")
		if ip == "" {
			ip = c.Request.Header.Get("X-Forwarded-For")
		}
		if ip == "" {
			ip = c.Request.RemoteAddr
		}
		// 字符串查找
		i := strings.Index(ip, ":")
		// 字符串截取
		ip = ip[0:i]
		if limiter.AllowIP(ip + ":" + c.Request.RequestURI) {
			fmt.Printf("Request from IP %s is allowed\n", ip)
			c.Next()
		} else {
			fmt.Printf("Request from IP %s is rate limited\n", ip)
			c.JSON(http.StatusTooManyRequests, tool.R{
				429, "请求被限流", nil,
			})
			c.Abort()
		}
	}
}
